# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Helpers for creating SRDataset.
"""

import os
from mindspore import context
from mindspore.context import ParallelMode
from mindspore.communication.management import init, get_rank, get_group_size
import mindspore.dataset as ds
import math
from .dataset import SRDataset


class MySampler():
    """sampler for distribution"""

    def __init__(self, dataset, local_rank, world_size):
        self.__num_data = len(dataset)
        self.__local_rank = local_rank
        self.__world_size = world_size
        self.samples_per_rank = int(math.ceil(self.__num_data / float(self.__world_size)))
        self.total_num_samples = self.samples_per_rank * self.__world_size

    def __iter__(self):
        """"iter"""
        indices = list(range(self.__num_data))
        indices.extend(indices[:self.total_num_samples - len(indices)])
        indices = indices[self.__local_rank:self.total_num_samples:self.__world_size]
        return iter(indices)

    def __len__(self):
        """length"""
        return self.samples_per_rank


def create_traindataset(dataset_path, scale, repeat_num=1,
                        batch_size=8, distribute=False):
    """
    Create an SRDataset for training or testing.

    Args:
        dataset_path (string): Path to the dataset.
        scale (int): downscaling ratio.
        do_train (bool): Whether dataset is used for training or testing.
        repeat_num (int): Repeat times of the dataset.
        batch_size (int): Batch size of the dataset.
        target (str): Device target.
        distribute (bool): For distributed training or not.

    Returns:
        dataset
    """
    paths = []
    for p, _, fs in sorted(os.walk(dataset_path)):
        for f in sorted(fs):
            if f.endswith(".png"):
                paths.append(os.path.join(p, f))
    assert paths, "no png images found"

    sr_ds = SRDataset(paths, scale=scale, training=True)

    parallel_mode = context.get_auto_parallel_context("parallel_mode")

    if distribute:
        rank_id = get_rank()
        rank_size = get_group_size()
    else:
        rank_size = 1

    num_shards = None if rank_size == 1 else rank_size
    shard_id = None if rank_size == 1 else rank_id

    if parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
        sampler = MySampler(sr_ds, local_rank=rank_id, world_size=rank_size)
        dataset = ds.GeneratorDataset(sr_ds, ["HR", "LR"], shuffle=True,
                                      sampler=sampler, num_shards=num_shards, shard_id=shard_id)
    else:
        dataset = ds.GeneratorDataset(sr_ds, ["HR", "LR"], shuffle=True)

    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.repeat(repeat_num)
    return dataset

def create_testdataset(dataset_path, scale, repeat_num=1,
                        batch_size=1, distribute=False):
    """
    Create an SRDataset for training or testing.

    Args:
        dataset_path (string): Path to the dataset.
        scale (int): downscaling ratio.
        do_train (bool): Whether dataset is used for training or testing.
        repeat_num (int): Repeat times of the dataset.
        batch_size (int): Batch size of the dataset.
        target (str): Device target.
        distribute (bool): For distributed training or not.

    Returns:
        dataset
    """
    paths = []
    for p, _, fs in sorted(os.walk(dataset_path)):
        for f in sorted(fs):
            if f.endswith(".png"):
                paths.append(os.path.join(p, f))
    assert paths, "no png images found"

    sr_ds = SRDataset(paths, scale=scale, training=False)

    parallel_mode = context.get_auto_parallel_context("parallel_mode")

    if distribute:
        rank_id = get_rank()
        rank_size = get_group_size()
    else:
        rank_size = 1

    num_shards = None if rank_size == 1 else rank_size
    shard_id = None if rank_size == 1 else rank_id

    if parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
        sampler = MySampler(sr_ds, local_rank=rank_id, world_size=rank_size)
        dataset = ds.GeneratorDataset(sr_ds, ["HR", "LR"], shuffle=True,
                                      sampler=sampler, num_shards=num_shards, shard_id=shard_id)
    else:
        dataset = ds.GeneratorDataset(sr_ds, ["HR", "LR"], shuffle=True)

    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.repeat(repeat_num)
    return dataset